import time
import torch
import torch.nn.functional as F

from dataset.metrics import AverageMeter, HeadTailAcc
from dataset.metrics import accuracy as accuracy_mine


def accuracy(y_pred, y_actual, topk=(1, ), return_tensor=False):
    """
    Computes the precision@k for the specified values of k in this mini-batch
    :param y_pred   : tensor, shape -> (batch_size, n_classes)
    :param y_actual : tensor, shape -> (batch_size)
    :param topk     : tuple
    :param return_tensor : bool, whether to return a tensor or a scalar
    :return:
        list, each element is a tensor with shape torch.Size([])
    """
    maxk = max(topk)
    batch_size = y_actual.size(0)

    _, pred = y_pred.topk(maxk, 1, True, True)
    pred = pred.t()
    correct = pred.eq(y_actual.view(1, -1).expand_as(pred))

    res = []
    for k in topk:
        correct_k = correct[:k].view(-1).float().sum(0)
        if return_tensor:
            res.append(correct_k.mul_(100.0 / batch_size))
        else:
            res.append(correct_k.item() * 100.0 / batch_size)
    return res


@torch.no_grad()
def evaluate(model, dataloader, logger):
    model.eval()
    end = time.time()
    losses = AverageMeter('Loss')
    top1 = AverageMeter('Acc@1')
    top5 = AverageMeter('Acc@5')

    num_classes = dataloader.dataset.nb_classes

    if num_classes in [10, 100]:
        head_tail_acc = HeadTailAcc(num_classes)
    else:
        head_tail_acc = None

    for sample in dataloader:
        x = sample['data'].cuda()
        labels = sample['label'].cuda()
        output = model(x)
        logits = output['logits']
        loss = F.cross_entropy(logits, labels)

        acc1, acc5 = accuracy_mine(logits.cpu(), labels.cpu(), topk=(1, 5))
        batch_size = x.size(0)
        top1.update(acc1[0], batch_size)
        top5.update(acc5[0], batch_size)
        losses.update(loss.item(), batch_size)

        if head_tail_acc:
            head_tail_acc.update_per_cls_acc(logits.cpu(), labels.cpu())

    logger.info(f'Evaluate Summary '
                f'Time {time.time() - end:.2f}s\t'
                f'{losses}\t {top1}\t {top5}')
    if head_tail_acc:
        logger.info(f'{head_tail_acc}')
    return {'accuracy': top1.avg, 'loss': losses.avg}